package matchbox;

import inference.GraphFactor;
import inference.GraphNode;
import inference.GraphUtils;
import inference.GraphVariable;
import inference.BeliefMessage;
import probability.GaussianDistribution;
import probability.ProbabilityDistribution;
import utils.Printer;
import basics.DenseMatrix;
import basics.DenseVector;
import basics.Matrix;
import basics.Vector;
import basics.VectorMatrix;

public class ZFactor extends MatchboxFactor {

	private GaussianDistribution _toS = new GaussianDistribution(0, 1);
	private GaussianDistribution _toT = new GaussianDistribution(0, 1);
	private GaussianDistribution _toZ = new GaussianDistribution(0, 1);
	private GaussianDistribution _zToThis;
	private GaussianDistribution _fromS;
	private GaussianDistribution _fromT;

	public ZFactor(String vname) {
		super(vname, Type.Operation);
		// _from = new GaussianDistribution(0, 1);
	}

	private GraphVariable getGraphVar(String name) {
		return (GraphVariable) edges(name).nodes().get(0);
	}

	@Override
	public GaussianDistribution ep(GraphNode next) {
		GaussianDistribution g = super.ep(next);
		if (next.is(NodeNames.S)) {
			_toS = g;
		} else if (next.is(NodeNames.T)) {
			_toT = g;
		} else if (next.is(NodeNames.Z)) {
			_toZ = g;
		}

		return g;
	}

	@Override
	public BeliefMessage toCached(GraphNode next) {
		if (next.is(NodeNames.S)) {
			return new BeliefMessage(_toS);
		} else if (next.is(NodeNames.T)) {
			return new BeliefMessage(_toT);
		} else { // if Z
			return new BeliefMessage(_toZ);
		}
	}

	@Override
	public BeliefMessage to(GraphNode next) {
		return to(next, null);
	}

	public BeliefMessage to(GraphNode next, String opt) {
		BeliefMessage m = null;
		if (next.is(NodeNames.S) || next.is(NodeNames.T)) { // *>s,*>t
			// if (opt == null) {
			// m = messageToSTSampling(next, opt);
			// }
			if (opt == null) {
				// m = messageToSTSampling(next, opt);
				m = messageToST(next);
			} else {
				_zToThis = (GaussianDistribution) getGraphVar(NodeNames.Z).to(this).getDistribution();

				if (opt.equals(NodeNames.S)) {
					m = messageToZfreezing(_zToThis, _fromS.getMu());
					_toT = (GaussianDistribution) m.getDistribution();
				} else {
					m = messageToZfreezing(_zToThis, _fromT.getMu());
					_toS = (GaussianDistribution) m.getDistribution();
				}
			}

			_fromS = _toS;
			_fromT = _toT;
		} else { // *>z
			GraphVariable svar = getGraphVar(NodeNames.S);
			GraphVariable tvar = getGraphVar(NodeNames.T);

			if (opt == null) {
				_fromS = (GaussianDistribution) svar.to(this).getDistribution();
				_fromT = (GaussianDistribution) tvar.to(this).getDistribution();

				m = messageToZExact(next);
				// m = messageToZVariational(next);

			} else if (opt.equals(NodeNames.S)) { // let's freeze s
				_fromT = (GaussianDistribution) tvar.to(this).getDistribution();
				_fromS = (GaussianDistribution) svar.toCached(this).getDistribution();
				m = messageToZfreezing(_fromT, _fromS.getMu()); // ((GaussianDistribution)
																// tvar.toCached(this).getDistribution()).getMu()

			} else if (opt.equals(NodeNames.T)) { // let's freeze t
				_fromS = (GaussianDistribution) svar.to(this).getDistribution();
				_fromT = (GaussianDistribution) tvar.toCached(this).getDistribution();
				m = messageToZfreezing(_fromS, _fromT.getMu()); // ((GaussianDistribution)
																// svar.toCached(this).getDistribution()).getMu()
			}

			_toZ = (GaussianDistribution) m.getDistribution();
		}

		Printer.print("F_Z>" + next + ":" + m.getDistribution().toString());

		return m; // new DenseVector(new double[] { _from.measure(m) });
	}

	public BeliefMessage toCached(GraphNode next, String opt) {
		BeliefMessage m = null;
		// || next.is(NodeNames.T)
		if (next.is(NodeNames.S)) { // *>s,*>t
			m = new BeliefMessage(_toS);
		} else if (next.is(NodeNames.T)) { // *>s,*>t
			m = new BeliefMessage(_toT);
		} else { // *>z
			m = new BeliefMessage(_toZ);
		}

		return m;
	}

	private BeliefMessage messageToZfreezing(GaussianDistribution from, double mu) {
		return new BeliefMessage(new GaussianDistribution(from.getMu() * mu, mu * mu * from.getSigma2()));
	}

	// This calculates the exact value for the outgoing messsage. Its mean is mu1*mu2 and var is
	// mu1^2sigma2^2+mu2^2sigma1^2+sigma1^2sigma2^2
	private BeliefMessage messageToZExact(GraphNode next) {
		double var = (Math.pow(_fromS.getMu(), 2) * _fromT.getSigma2())
				+ (Math.pow(_fromT.getMu(), 2) * _fromS.getSigma2()) + (_fromS.getSigma2() * _fromT.getSigma2());
		return new BeliefMessage(new GaussianDistribution(_fromS.getMu() * _fromT.getMu(), var));
	}

	private BeliefMessage messageToZSampling(GraphNode next) {
		_toZ = GaussianDistribution.sampleIndicatorProduct(5, _fromS, _fromT);
		return new BeliefMessage(_toZ);
	}

	private BeliefMessage messageToSTSampling(GraphNode next, String freezedVar) {
		_zToThis = (GaussianDistribution) getGraphVar(NodeNames.Z).to(this).getDistribution();

		updateSTSampling(freezedVar);

		GraphVariable s = getGraphVar(NodeNames.S);
		if (next == s) {
			return new BeliefMessage(_toS);
		} else {
			return new BeliefMessage(_toT);
		}
	}

	private void updateSTSampling(String freezedVar) {
		// for (int i = 0; i < 5; i++) {
		if (freezedVar == null || freezedVar.equals(NodeNames.T)) {
			_toS = GaussianDistribution.sampleIndicatorDivide(5, _zToThis, _fromT);
		}
		if (freezedVar == null || freezedVar.equals(NodeNames.S)) {
			_toT = GaussianDistribution.sampleIndicatorDivide(5, _zToThis, _fromS);
		}
		// }
	}

	private BeliefMessage messageToST(GraphNode next) {
		_zToThis = (GaussianDistribution) getGraphVar(NodeNames.Z).to(this).getDistribution();
		// updateST((GraphVariable) next);
		updateST(getGraphVar(NodeNames.S));
		updateST(getGraphVar(NodeNames.T));
		
		if (next.is(NodeNames.S))
			return new BeliefMessage(_toS);

		return new BeliefMessage(_toT);
	}

	private static GaussianDistribution getPst(GaussianDistribution toST, GaussianDistribution fromST) {
		return new GaussianDistribution(toST.getMu() + fromST.getMu(), toST.getSigma2() + fromST.getSigma2());
	}

	public GaussianDistribution updateST(GraphVariable st) {
		GraphVariable s = getGraphVar(NodeNames.S);

		GaussianDistribution other;
		GaussianDistribution g;

		if (st == s) {
			other = getPst(_toT, _fromT);
			g = _toS;
		} else {
			other = getPst(_toS, _fromS);
			g = _toT;
		}

		g.setMu(_zToThis.getMu() * other.getMu() / other.getRawMoment(2));
		g.setSigma2(_zToThis.getRawMoment(2) - Math.pow(_zToThis.getMu(), 2) / other.getRawMoment(2));
		g.setSigma2(g.getSigma2() / 10.);

		return g;
	}

	private BeliefMessage messageToZVariational(GraphNode next) {
		// GraphVariable svar = getGraphVar(NodeNames.S);
		// GraphVariable tvar = getGraphVar(NodeNames.T);
		//
		// _fromS = (GaussianDistribution) svar.to(this).getDistribution();
		// _fromT = (GaussianDistribution) tvar.to(this).getDistribution();

		GaussianDistribution s = getPst(_toS, _fromS);
		GaussianDistribution t = getPst(_toT, _fromT);

		_toZ.setMu(s.getMu() * t.getMu());
		_toZ.setSigma2(s.getRawMoment(2) * t.getRawMoment(2) - Math.pow(s.getMu(), 2) * Math.pow(t.getMu(), 2));
		_toZ.setSigma2(_toZ.getSigma2() / 3.);

		return new BeliefMessage(_toZ); // , (Vector) fromS.getValue().timesCopy(fromT.getValue())
	}
	// private VectorMatrix stMean(Message st) {
	// return st.getDistribution().getParam(0).addCopy(_from.getParam(0));
	// }
	//
	// private VectorMatrix stVar(Message st) {
	// return st.getDistribution().getParam(1).addCopy(_from.getParam(1));
	// }

}
